import json
import re
from datetime import datetime, timedelta
from typing import Dict, List
from fastapi import FastAPI
from pydantic import BaseModel

from langchain.tools.render import render_text_description
from langchain_core.messages import AIMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

url_base = "http://localhost:11434/v1"
api_key = "11"
qwen_llm = ChatOpenAI(openai_api_base=url_base, openai_api_key=api_key, model="qwen2")

# 添加一个简单的数据存储（实际项目中应该使用数据库）
meeting_records: Dict[str, List[Dict]] = {}


def parse_chinese_time(time_str: str) -> datetime:
    """
    解析中文时间表达式
    支持格式：
    - 明天下午3点
    - 后天上午10点
    - 下周一下午2点
    - 具体日期时间(2024-03-20 14:00)
    """
    now = datetime.now()

    # 首先尝试直接解析标准格式
    try:
        return datetime.strptime(time_str, '%Y-%m-%d %H:%M')
    except ValueError:
        pass

    # 解析中文时间表达式
    day_mapping = {
        "今天": 0,
        "明天": 1,
        "后天": 2,
        "大后天": 3
    }

    time_pattern = {
        "上午": (9, 12),
        "中午": (12, 13),
        "下午": (13, 18),
        "晚上": (18, 22)
    }

    # 提取日期信息
    target_date = now
    for day, offset in day_mapping.items():
        if day in time_str:
            target_date = now + timedelta(days=offset)
            break

    # 提取具体时间
    hour = 0
    minute = 0

    # 解析"下午3点"这样的格式
    for period, (start_hour, end_hour) in time_pattern.items():
        if period in time_str:
            # 提取数字
            numbers = re.findall(r'\d+', time_str)
            if numbers:
                hour = int(numbers[0])
                if period in ["下午", "晚上"] and hour < 12:
                    hour += 12
                break

    # 提取分钟
    minute_match = re.search(r'[:：](\d{1,2})', time_str)
    if minute_match:
        minute = int(minute_match.group(1))

    return target_date.replace(hour=hour, minute=minute)


@tool
def check_meeting(date: str, room_no: str, user_name: str) -> dict:
    """
    检查会议室是否可用
    Args:
        date: 日期时间，支持格式：
            - 具体时间：'2024-03-20 14:00'
            - 相对时间：'明天下午3点'
        room_no: 会议室编号
        user_name: 预订人姓名
    Returns:
        dict: 包含可用状态和消息的字典
    """
    try:
        # 解析时间
        meeting_time = parse_chinese_time(date)

        # 格式化时间为标准格式
        formatted_date = meeting_time.strftime('%Y-%m-%d %H:%M')

        # 获取该会议室的预订记录
        room_bookings = meeting_records.get(room_no, [])

        # 检查是否有时间冲突
        for booking in room_bookings:
            booked_time = datetime.strptime(booking['date'], '%Y-%m-%d %H:%M')
            # 假设每个会议持续1小时，检查是否有重叠
            if abs((meeting_time - booked_time).total_seconds()) < 3600:
                return {
                    "available": False,
                    "message": f"会议室已被{booking['user_name']}预订",
                    "conflict_time": booking['date']
                }

        # 如果没有冲突，添加新预订
        new_booking = {
            "date": formatted_date,
            "user_name": user_name
        }
        if room_no not in meeting_records:
            meeting_records[room_no] = []
        meeting_records[room_no].append(new_booking)

        return {
            "available": True,
            "message": "会议室预订成功",
            "booking_info": new_booking
        }

    except ValueError as e:
        return {
            "available": False,
            "message": f"日期格式错误: {str(e)}"
        }


@tool
def add(first_int: int, second_int: int):
    "Add two integers."
    return first_int + second_int


@tool
def exponentiate(base: int, exponent: int) -> int:
    "Exponentiate the base to the exponent power."
    return base ** exponent


@tool
def multiply(first_int: int, second_int: int) -> int:
    """Multiply two integers together."""
    return first_int * second_int


def extract_format_json(message: AIMessage):
    """Extracts JSON content from a string where JSON is embedded between ```json and ``` tags.
    """

    text = message.content
    try:
        return json.loads(text)
    except Exception:
        raise ValueError(f"Failed to parse: {message}")


tools = [add, multiply, exponentiate, check_meeting]


def tool_chain(model_output):
    try:
        tool_map = {tool.name: tool for tool in tools}

        # 检查输出是否为字典类型
        if not isinstance(model_output, dict):
            return {"response": str(model_output)}

        chosen_tool = tool_map.get(model_output.get("name"))

        if chosen_tool:
            # 使用 invoke 方法替代直接调用
            return chosen_tool.invoke(model_output["arguments"])
        else:
            # 如果不是工具调用，直接返回模型的回答
            return model_output.get("arguments", {}).get("response", model_output)
    except Exception as e:
        print(f"处理工具调用时出错: {e}")
        return {"error": str(e)}


rendered_tools = render_text_description(tools)

system_prompt = """You are an assistant that has access to the following set of tools. Here are the names and descriptions for each tool:

{rendered_tools}

If the user's question can be answered using one of the tools, return the name and input of the tool to use in JSON format.
If the question cannot be answered using tools, provide a direct response in this JSON format:
{{"name": "response", "arguments": {{"response": "你的回答"}}}}

请确保输出是有效的 JSON 格式，并包含 'name' 和 'arguments' 两个键。"""

sample = """
参考示例
add: add(first_int: int, second_int: int) -> int - Add two integers.
输出示例
{{"name":"add","arguments":{{"first_int":...,"second_int":...}}}}
"""

prompt = ChatPromptTemplate.from_messages(
    [("system", system_prompt), ("system", sample), ("user", "{input}")]
)

# 修改这里：创建一个包含所有必需变量的链
chain = (
        {"input": RunnablePassthrough(), "rendered_tools": lambda _: rendered_tools}
        | prompt
        | qwen_llm
        | extract_format_json
        | RunnablePassthrough.assign(output=tool_chain)
)

# 创建 FastAPI 应用
app = FastAPI()


# 定义请求体模型
class MeetingRequest(BaseModel):
    input: str


@app.post("/invoke")
def invoke_tool_chain(request: MeetingRequest):
    """调用工具链并返回结果"""
    result = chain.invoke({"input": request.input})
    return result


# 运行 FastAPI 应用
# 在终端中运行: uvicorn demo:app --reload

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
